package com.lagou.hdfs.test_2;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.GzipCodec;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.TextOutputFormat;
import org.apache.hadoop.util.ReflectionUtils;

import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;

/**
 * @ClassName MyOutFormat
 * @Description TODO
 * @Author huangjh5
 * @Date 2020/12/9 15:48
 */
public class MyOutFormat<K, V> extends TextOutputFormat<K, V> {


    public MyOutFormat() {
    }



    public RecordWriter<K, V> getRecordWriter(TaskAttemptContext job) throws IOException, InterruptedException {
        Configuration conf = job.getConfiguration();
        boolean isCompressed = getCompressOutput(job);
        String keyValueSeparator = conf.get(SEPERATOR, "\t");
        CompressionCodec codec = null;
        String extension = "";
        if (isCompressed) {
            Class<? extends CompressionCodec> codecClass = getOutputCompressorClass(job, GzipCodec.class);
            codec = (CompressionCodec) ReflectionUtils.newInstance(codecClass, conf);
            extension = codec.getDefaultExtension();
        }

        Path file = this.getDefaultWorkFile(job, extension);
        FileSystem fs = file.getFileSystem(conf);
        FSDataOutputStream fileOut = fs.create(file, false);
        return isCompressed ? new MyOutFormat.MyLineRecordWriter(new DataOutputStream(codec.createOutputStream(fileOut)), keyValueSeparator) : new MyOutFormat.MyLineRecordWriter(fileOut, keyValueSeparator);
    }

    static class  MyLineRecordWriter extends LineRecordWriter{

        //统计条数
        private Integer recordNum = 1;


        private final byte[] myKeyValueSeparator;

        public MyLineRecordWriter(DataOutputStream out, String keyValueSeparator) {
            super(out, keyValueSeparator);
            this.myKeyValueSeparator = keyValueSeparator.getBytes(StandardCharsets.UTF_8);
        }


        @Override
        public synchronized void write(Object key, Object value) throws IOException {
            this.out.write(recordNum.toString().getBytes(StandardCharsets.UTF_8));
            recordNum ++;
            this.out.write(this.myKeyValueSeparator);
            super.write(key, value);
        }
    }

}
